'''
Created on Jun 20, 2012

@author: alebalbin
'''

import sys
import numpy as np
import array
from optparse import OptionParser
from collections import defaultdict,deque #, OrderedDict
from bx.intervals.intersection import Interval, IntervalTree
from lib.parse_gtf2 import read_gtf

JOB_SUCCESS=0
JOB_ERROR=1

'''
class gtf_transcript():
    def __init__(self,fields):
        self.annotation=defaultdict()
        self.chrom=fields[0]
        self.feature=fields[1] 
        self.feature_type=fields[2]
        self.start=fields[3]
        self.end=fields[4]
        self.atributte1=fields[5]
        self.strand=fields[6]
        self.atribute2=fields[7]
        for a in fields[-1].split(';'):
            if a=='':
                continue
            att, value= a.split()[0].strip(),a.split()[1].strip()
            self.annotation[att]=value
        self.returnme()
    def returnme(self):
        return self
'''

class transcript():
    
    def __init__(self):
        self.id=None
        self.chr=0
        self.start=0
        self.end=0
        self.latest_exon=0
        self.gtftrans=0
        self.FPKM=defaultdict()
        
    
    def grow_transcript(self,gtf_transcript):
        #compare start/end of this new transcript 
        # against the stored start/end
        if "exon_number" not in  gtf_transcript.annotation.keys():
            thisexon=1
        else:
            thisexon = int(gtf_transcript.annotation["exon_number"])
        start=gtf_transcript.start
        end=gtf_transcript.end
        
        if self.id is not None:            
            if self.latest_exon < thisexon:
                # check new end coordinate
                if end > self.end:
                    self.end=end
                    self.FPKM[thisexon]=gtf_transcript.annotation["FPKM"]
            if self.latest_exon > thisexon:
                # check new start coordinate
                if start < self.start:
                    self.start=start
                    self.FPKM[thisexon]=gtf_transcript.annotation["FPKM"]
        else:
            self.id=gtf_transcript.annotation["transcript_id"]
            self.chrom=gtf_transcript.chrom
            self.feature=gtf_transcript.feature
            self.feature_type=gtf_transcript.feature_type
            self.start=gtf_transcript.start
            self.end=gtf_transcript.end
            self.atributte1=gtf_transcript.atributte1
            self.strand=gtf_transcript.strand
            self.atribute2=gtf_transcript.atribute2
            self.latest_exon=thisexon
            self.gtftrans=gtf_transcript
            self.annotation=gtf_transcript.annotation.copy()
            self.FPKM[thisexon]=gtf_transcript.annotation["FPKM"]



        # this will assign as the gtf_info the latest 
        # gtf_transcript info. From cufflinks it seems
        # that all annotation is the same for all exons in the transcript.
        # this seems supicious as each exon could have a different FPKM.
        self.gtftrans=gtf_transcript
        
        
class transcriptome():
    
    def __init__(self):
        self.transcripts_list=defaultdict(transcript)
    
    def grow_transcriptome(self,t):
        '''
        t = gtf_transcript class
        '''
        self.transcripts_list[t.annotation["transcript_id"]].grow_transcript( t )
                

class strandDNATree():
    
    def __init__(self):
        self.genome=defaultdict()
        '''
        chrlist = ['chr1', 'chr2','chr3','chr4','chr5','chr6','chr7', 'chr8','chr9','chr10', 'chr11','chr12','chr6','chr7', 'chr8','chr9','chr10', 'chr11', 'chr12', \
                   'chr13', 'chr14','chr15','chr16','chr17','chr18','chr19', 'chr20','chr21','chr22', 'chrX','chrY']
        '''
        self.maxPosStrand=0
        self.maxNegStrand=0
        strands=['+','-']
        for s in strands:
            self.genome[s]=IntervalTree()
    
    def grow_strand(self,t):
        '''
        t=transcript
        '''
        # Checks that the transcript have the minimun information required
        #
        # Chcek the strand
        if t is not None:
            if t.strand=="+":
                self.genome['+'].insert_interval( Interval(t.start, t.end, t.annotation, 
                                                           t.chrom, t.strand) )
                # get the maximum end point in the positive strand
                if t.end > self.maxPosStrand:
                    self.maxPosStrand=t.end
                
            elif t.strand=="-":
                self.genome['-'].insert_interval( Interval(t.start, t.end, t.annotation, 
                                                           t.chrom, t.strand) )
                # get the maximum end point in the positive strand
                if t.end > self.maxNegStrand:
                    self.maxNegStrand=t.end

            #print t.feature,t.annotation['transcript_id'],t.start,t.end, t.strand

        else:
            print "Excluded Transcript"
    
    def classify_overlapping_trancript(self,intv, start, end):
        '''
        intv is the interval intersected from querying the interval tree 
        '''
        if (intv.start >= start) and (intv.end <= end):
            overlap_bp=intv.end-intv.start
            return "%s_%d"%("OP",overlap_bp)#"OVERLAPPING"
        elif (start <= intv.start) and (end <= intv.end):
            overlap_bp=end-intv.start
            return "%s_%d"%("RL",overlap_bp)# "RL"#"TAIL_TO_TAIL"
        elif (intv.start <= start) and (end <= intv.end):
            overlap_bp=end-start
            return "%s_%d"%("OP",overlap_bp)#"OP"
        elif (intv.start <= start) and (end >= intv.end):
            overlap_bp=intv.end-start
            return "%s_%d"%("LR",overlap_bp)# "LR"#"HEAD_TO_HEAD"
        else:
            return "UN"#"Not determined"
            
            
    def sense_nat_pair(self,t,strand):
        '''
        '''
        intervals=self.genome[strand].find(t.start, t.end)
        if intervals:
            for intv in intervals:
                if intv.chrom == t.chrom:
                    #print t.start, t.end, t.strand, intv.start,intv.chrom,intv.strand
                    # note gene_name key in ensemble, gene_id key in cufflinks
                    # FPKM only exists in cufflinks.
                    overlap_type=self.classify_overlapping_trancript(intv,t.start, t.end)
                    original_transc=[t.value['transcript_id'],t.value['gene_id'],t.chrom,t.start, t.end, t.strand,t.value['FPKM']]
                    overlap_transc=[intv.value['transcript_id'],intv.value['gene_id'], intv.chrom, intv.start, intv.end, intv.strand,intv.value['FPKM']]
                    #print original_transc, overlap_transc
                    # 0.5 added to avoid division by zero.
                    return original_transc + overlap_transc + [overlap_type]+[(np.log2(float(t.value['FPKM'])+0.5) - np.log2(float(intv.value['FPKM'])+0.5))]
        

